{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# AI Assistant Extensibility\n",
    "\n",
    "This notebook presents how to extend the AI Assistant plugin in the runtime by providing a custom message handler."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "%use intellij-platform",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Load the AI Assistant bundled plugin using its `com.intellij.ml.llm` ID."
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "loadPlugins(\"com.intellij.ml.llm\")",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import com.intellij.ml.llm.agents.ChatAgent\n",
    "import com.intellij.ml.llm.core.chat.messages.ChatMessage\n",
    "import com.intellij.ml.llm.core.chat.messages.UserMessage\n",
    "import com.intellij.ml.llm.core.chat.messages.impl.SimpleCompletableMessage\n",
    "import com.intellij.ml.llm.core.chat.session.ChatKind\n",
    "import com.intellij.ml.llm.core.chat.session.ChatMessageHandler\n",
    "import com.intellij.ml.llm.core.chat.session.ChatSession\n",
    "import com.intellij.ml.llm.smartChat.endpoints.SmartChatEndpoint\n",
    "import com.intellij.openapi.application.ApplicationManager\n",
    "import com.intellij.openapi.extensions.ExtensionPointName\n",
    "import com.intellij.openapi.project.Project\n",
    "\n",
    "class MyChatMessageHandler : ChatMessageHandler {\n",
    "    override fun isApplicable(project: Project, kind: ChatKind, userMessage: UserMessage) = true\n",
    "\n",
    "    override fun createAnswerMessage(\n",
    "        project: Project,\n",
    "        chat: ChatSession,\n",
    "        userMessage: UserMessage,\n",
    "        kind: ChatKind,\n",
    "    ) = SimpleCompletableMessage(chat)\n",
    "\n",
    "    override suspend fun serveAnswerMessage(\n",
    "        project: Project,\n",
    "        chat: ChatSession,\n",
    "        answerMessage: ChatMessage,\n",
    "        smartChatEndpoints: List<SmartChatEndpoint>,\n",
    "    ) {\n",
    "        ChatAgent.EP_NAME.extensionList\n",
    "            .find { it.id == \"groot\" }\n",
    "            ?.serveAnswerMessage(project, chat, answerMessage)\n",
    "    }\n",
    "}\n",
    "\n",
    "registerExtension(ChatMessageHandler.EP, MyChatMessageHandler())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import com.intellij.ml.llm.agents.ChatAgent\n",
    "import com.intellij.ml.llm.core.chat.messages.ChatMessage\n",
    "import com.intellij.ml.llm.core.chat.messages.impl.SimpleCompletableMessage\n",
    "import com.intellij.ml.llm.core.chat.session.ChatSession\n",
    "import com.intellij.ml.llm.privacy.trustedStringBuilders.privacyConst\n",
    "import com.intellij.openapi.project.Project\n",
    "\n",
    "class MyChatAgent : ChatAgent {\n",
    "\n",
    "    override val id = \"groot\"\n",
    "    override val name = \"I am Groot\"\n",
    "\n",
    "    override fun createAnswerMessage(\n",
    "        project: Project,\n",
    "        chat: ChatSession,\n",
    "        userMessage: UserMessage,\n",
    "        kind: ChatKind,\n",
    "    ) = SimpleCompletableMessage(chat)\n",
    "\n",
    "    override suspend fun serveAnswerMessage(\n",
    "        project: Project,\n",
    "        chat: ChatSession,\n",
    "        answerMessage: ChatMessage,\n",
    "    ) {\n",
    "        answerMessage as SimpleCompletableMessage\n",
    "        answerMessage.appendText(\"*I am Groot*\".privacyConst)\n",
    "    }\n",
    "}\n",
    "\n",
    "registerExtension(ChatAgent.EP_NAME, MyChatAgent())"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "name": "kotlin",
   "version": "1.9.23",
   "mimetype": "text/x-kotlin",
   "file_extension": ".kt",
   "pygments_lexer": "kotlin",
   "codemirror_mode": "text/x-kotlin",
   "nbconvert_exporter": ""
  },
  "ktnbPluginMetadata": {
   "sessionRunMode": "IDE_PROCESS"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}